﻿using System;
using System.Runtime.InteropServices;
using Microsoft.Win32;
using System.Collections.Generic;
using System.Reflection;
using System.IO;

namespace Tsanie.ShellExt
{

    #region - Custom Attribute -

    [AttributeUsage(AttributeTargets.Class, AllowMultiple = false, Inherited = false)]
    public sealed class ExtensionFileTypesAttribute : Attribute
    {
        private bool m_findAssoc = false;
        private string[] m_FileTypes;

        public ExtensionFileTypesAttribute(string fileTypes) { m_FileTypes = fileTypes.Split(','); }
        public string FileTypes { get { return String.Join(",", m_FileTypes); } }
        public string[] FileTypesArray { get { return m_FileTypes; } }
        public bool FindAssoc
        {
            get { return m_findAssoc; }
            set { m_findAssoc = value; }
        }
    }

    #endregion

    delegate bool RegDelegate(string assembly, string className);

    public class Common
    {
        public const string CONFIG_NAME = "ShellExtManager.txt";
        public const string CONFIG_MENUTYPE = "Tsanie.ShellExt.ManagerContextMenu";
        public const string CONFIG_PROPERTYTYPE = "Tsanie.ShellExt.ManagerPropertySheet";

        public const string GUID_CONTEXTMENU = "{CE80D400-D942-4730-92B8-074E9EF12AD8}";
        public const string GUID_THUMBNAIL = "{B0D9EBC0-5612-4359-B91B-BECA3365CA45}";
        public const string GUID_INFOTIP = "{DFB16C44-4FEB-4432-938F-D0318CAF155B}";
        public const string GUID_PROPERTYSHEET = "{D2367552-6A58-4621-9A56-2C549E8D215D}";

        #region - APIs -

        [DllImport("kernel32.dll", CharSet = CharSet.Auto)]
        public static extern void OutputDebugString(string s);
        [DllImport("Shell32.dll")]
        public static extern void SHChangeNotify(int wEventId, uint uFlags, IntPtr dwItem1, IntPtr dwItem2);
        public const int SHCNE_ASSOCCHANGED = 0x08000000;
        [DllImport("uxtheme.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)]
        public extern static int SetWindowTheme(IntPtr hWnd, string pszSubAppName, string pszSubIdList);

        #endregion

        #region - Debug Method -

        public static void OutputDebugString(string format, object arg1)
        {
            OutputDebugString(format, new object[] { arg1 });
        }
        public static void OutputDebugString(string format, object arg1, object arg2)
        {
            OutputDebugString(format, new object[] { arg1, arg2 });
        }
        public static void OutputDebugString(string format, object arg1, object arg2, object arg3)
        {
            OutputDebugString(format, new object[] { arg1, arg2, arg3 });
        }
        public static void OutputDebugString(string format, params object[] args)
        {
            OutputDebugString(string.Format(format, args));
        }

        public static void OutputError(string title, Exception e)
        {
            Common.OutputDebugString("[" + title + "] =======================");
            Common.OutputDebugString(e.GetType().FullName + ", " + e.Message);
            Common.OutputDebugString(e.StackTrace);
            Common.OutputDebugString("-----------------------------------------");
        }

        #endregion

        #region - Register/Unregister Method -

        static string CreateExtensionKey(string assoc, string extensionKey, string guid)
        {
            string result = assoc + "\\shellex\\" + extensionKey;
            try
            {
                using (RegistryKey key = Registry.ClassesRoot.CreateSubKey(result))
                {
                    key.SetValue(null, guid);
                    key.Close();
                }
                return result;
            }
            catch (Exception e)
            {
                Common.OutputDebugString("CreateExtensionKey():: Exception, '{0}', {1} - {2}", result, e.GetType().FullName, e.Message);
            }
            return null;
        }
        static string DeleteExtensionKey(string key)
        {
            try
            {
                Registry.ClassesRoot.DeleteSubKeyTree(key);
                return key;
            }
            catch (Exception e)
            {
                Common.OutputDebugString("DeleteExtensionKey(string):: Exception, '{0}', {1} - {2}", key, e.GetType().FullName, e.Message);
            }
            return null;
        }
        static string GetClassGuid(Type t)
        {
            GuidAttribute GuidAtt;
            try
            {
                GuidAtt = (GuidAttribute)(t.GetCustomAttributes(typeof(GuidAttribute), true)[0]);
                return "{" + GuidAtt.Value + "}";
            }
            catch (Exception e)
            {
                if (t == null)
                    Common.OutputDebugString("GetClassGuid(Type):: Argument 't' is null. - {1}", e.Message);
                else
                    Common.OutputDebugString("GetClassGuid(Type):: '{0}' don't have GuidAttribute. - {1}", t.FullName, e.Message);
            }
            return null;
        }

        internal static void RegisterExtension(Type t, string extensionKey)
        {
            string guid = GetClassGuid(t);
            if (string.IsNullOrEmpty(guid))
            {
                Common.OutputDebugString("RegisterExtension(Type, string):: [register] failed. Guid of '{0}' is null or empty.", t.FullName);
                return;
            }
            using (RegistryKey key = Registry.LocalMachine.CreateSubKey(@"Software\Microsoft\Windows\CurrentVersion\Shell Extensions\Approved"))
            {
                key.SetValue(guid, t.FullName);
                key.Close();
                //Common.OutputDebugString(@"RegisterExtension(Type, string):: Register extension approved - 'HKLM\Software\Microsoft\Windows\CurrentVersion\Shell Extensions\Approved, {0} = {1}'", guid, t.FullName);
            }
            try
            {
                object[] attrs = t.GetCustomAttributes(typeof(ExtensionFileTypesAttribute), true);
                if (attrs != null && attrs.Length > 0)
                {
                    ExtensionFileTypesAttribute fileTypes = (ExtensionFileTypesAttribute)attrs[0];
                    foreach (string type in fileTypes.FileTypesArray)
                    {
                        string result = null;
                        if (fileTypes.FindAssoc)
                        {
                            // 按照优先级依次查找关联
                            using (RegistryKey fileExt = Registry.CurrentUser.OpenSubKey(@"Software\Microsoft\Windows\CurrentVersion\Explorer\FileExts\" + type + @"\UserChoice"))
                            {
                                if (fileExt != null)
                                {
                                    string assocName = fileExt.GetValue("Progid") as string;
                                    fileExt.Close();
                                    if (assocName != null)
                                    {
                                        result = CreateExtensionKey(assocName, extensionKey, guid);
                                    }
                                }
                                if (result == null)
                                {
                                    // 查找HKCL\(ext)
                                    using (RegistryKey ext = Registry.ClassesRoot.OpenSubKey(type))
                                    {
                                        if (ext != null)
                                        {
                                            string assocName = ext.GetValue(null) as string;
                                            ext.Close();
                                            if (assocName != null)
                                            {
                                                result = CreateExtensionKey(assocName, extensionKey, guid);
                                            }
                                        }
                                    }
                                }
                            }
                        }
                        if (result == null)
                        {
                            // 如果都没有，可能是新建类型
                            result = CreateExtensionKey(type, extensionKey, guid);
                        }
                        //Common.OutputDebugString("RegisterExtension(Type, string):: [register] successfully. '{0}'", result);
                    }

                    // Tell Explorer to refresh
                    //Common.SHChangeNotify(Common.SHCNE_ASSOCCHANGED, 0, IntPtr.Zero, IntPtr.Zero);
                }
            }
            catch (Exception e)
            {
                Common.OutputDebugString("RegisterExtension(Type, string):: [register] exception: {0} - {1}", e.GetType().FullName, e.Message);
            }
        }
        internal static void UnregisterExtension(Type t, string extensionKey)
        {
            string guid = GetClassGuid(t);
            if (string.IsNullOrEmpty(guid))
            {
                Common.OutputDebugString("UnregisterExtension(Type, string):: [unregister] failed. Guid of '{0}' is null or empty.", t.FullName);
                return;
            }
            using (RegistryKey key = Registry.LocalMachine.CreateSubKey(@"Software\Microsoft\Windows\CurrentVersion\Shell Extensions\Approved"))
            {
                key.DeleteValue(guid);
                key.Close();
                //Common.OutputDebugString(@"UnregisterExtension(Type, string):: Unregister extension approved - 'HKLM\Software\Microsoft\Windows\CurrentVersion\Shell Extensions\Approved, {0} = {1}'", guid, t.FullName);
            }
            try
            {
                object[] attrs = t.GetCustomAttributes(typeof(ExtensionFileTypesAttribute), true);
                if (attrs != null && attrs.Length > 0)
                {
                    ExtensionFileTypesAttribute fileTypes = (ExtensionFileTypesAttribute)attrs[0];
                    foreach (string type in fileTypes.FileTypesArray)
                    {
                        string result = null;
                        if (fileTypes.FindAssoc)
                        {
                            // 按照优先级依次查找关联
                            using (RegistryKey fileExt = Registry.CurrentUser.OpenSubKey(@"Software\Microsoft\Windows\CurrentVersion\Explorer\FileExts\" + type + @"\UserChoice"))
                            {
                                if (fileExt != null)
                                {
                                    string assocName = fileExt.GetValue("Progid") as string;
                                    fileExt.Close();
                                    if (assocName != null)
                                    {
                                        result = DeleteExtensionKey(assocName + @"\shellex\" + extensionKey);
                                    }
                                }
                                if (result == null)
                                {
                                    // 查找HKCL\(ext)
                                    using (RegistryKey ext = Registry.ClassesRoot.OpenSubKey(type))
                                    {
                                        if (ext != null)
                                        {
                                            string assocName = ext.GetValue(null) as string;
                                            ext.Close();
                                            if (assocName != null)
                                            {
                                                result = DeleteExtensionKey(assocName + @"\shellex\" + extensionKey);
                                            }
                                        }
                                    }
                                }
                            }
                        }
                        if (result == null)
                        {
                            // 如果都没有，可能是新建类型
                            result = DeleteExtensionKey(type + @"\shellex\" + extensionKey);
                        }
                        //Common.OutputDebugString("UnregisterExtension(Type, string):: [unregister] successfully. '{0}'", result);
                    }

                    // Tell Explorer to refresh
                    Common.SHChangeNotify(Common.SHCNE_ASSOCCHANGED, 0, IntPtr.Zero, IntPtr.Zero);
                }
            }
            catch (Exception e)
            {
                Common.OutputDebugString("UnregisterExtension(Type, string):: [unregister] exception - {0}", e.Message);
            }
        }

        #endregion

        static List<BaseProvider> _providers;
        //static Assembly CurrentDomain_AssemblyResolve(object sender, ResolveEventArgs args)
        //{
        //    return Assembly.LoadFrom(Assembly.GetExecutingAssembly().Location);
        //}

        /// <summary>
        /// 获取当前加载的所有 Provider
        /// </summary>
        internal static List<BaseProvider> Providers
        {
            get
            {
                if (_providers == null)
                {
                    _providers = new List<BaseProvider>();
                    try
                    {
                        // 工作路径
                        string path = Assembly.GetExecutingAssembly().Location;
                        path = path.Substring(0, path.LastIndexOf('\\') + 1);
                        //AppDomainSetup ads = new AppDomainSetup
                        //{
                        //    ApplicationBase = path
                        //};
                        //AppDomain domain = AppDomain.CreateDomain("Tsanie.ShellExt.RegDomain", null, ads);
                        // .NET的一个bug，remote调用后在Unwrap时会提示缺少程序集，需要手动指定加载 ★
                        //AppDomain.CurrentDomain.AssemblyResolve += new ResolveEventHandler(CurrentDomain_AssemblyResolve);

                        // 插件配置文件
                        using (StreamReader reader = new StreamReader(path + CONFIG_NAME))
                        {
                            string line;
                            while ((line = reader.ReadLine()) != null)
                            {
                                if (line.Length > 0 && line[0] != '#')
                                {
                                    string[] kv = line.Split(',');
                                    if (kv.Length != 2)
                                    {
                                        continue;
                                    }
                                    try
                                    {
                                        //_providers.Add((BaseProvider)domain.CreateInstanceAndUnwrap(kv[0], kv[1]));
                                        _providers.Add((BaseProvider)(Assembly.Load(kv[0]).CreateInstance(kv[1])));
                                    }
                                    catch (Exception e)
                                    {
                                        Common.OutputError(string.Format("provider '{0}', '{1}'", kv[0], kv[1]), e);
                                    }
                                }
                            }
                            reader.Close();
                        }

                        //AppDomain.Unload(domain);
                    }
                    catch (Exception e)
                    {
                        Common.OutputError("get_Providers()", e);
                    }
                }
                return _providers;
            }
        }
        internal static bool RegisterKey(string ext, string regKey, ProviderType type)
        {
            bool changed = false;
            // 验证注册表是否缺少相关项
            string result = ext + @"\shellex\" + regKey;
            try
            {
                string guid;
                switch (type)
                {
                    case ProviderType.Thumbnail:
                        guid = GUID_THUMBNAIL;
                        break;
                    case ProviderType.ContextMenu:
                        guid = GUID_CONTEXTMENU;
                        // ContextMenu需要特殊处理
                        using (RegistryKey k = Registry.ClassesRoot.OpenSubKey(ext))
                        {
                            if (k != null)
                            {
                                result = k.GetValue(null) + @"\shellex\" + regKey;
                            }
                        }
                        break;
                    case ProviderType.InfoTip:
                        guid = GUID_INFOTIP;
                        break;
                    case ProviderType.PropertySheet:
                        guid = GUID_PROPERTYSHEET;
                        // PropertySheet的也需要特殊处理
                        using (RegistryKey k = Registry.ClassesRoot.OpenSubKey(ext))
                        {
                            if (k != null)
                            {
                                result = k.GetValue(null) + @"\shellex\" + regKey;
                            }
                        }
                        break;
                    default:
                        // 不支持的类型
                        return false;
                }
                using (RegistryKey k = Registry.ClassesRoot.OpenSubKey(result, true))
                {
                    if (k == null)
                    {
                        using (RegistryKey key = Registry.ClassesRoot.CreateSubKey(result))
                        {
#if DEBUG
                            Common.OutputDebugString("'{0}' - Create key.", result);
#endif
                            key.SetValue(null, guid);
                            changed = true;
                            key.Close();
                        }
                    }
                    else
                    {
                        if (changed || k.GetValue(null) as string != guid)
                        {
#if DEBUG
                            Common.OutputDebugString("'{0}' - Has key, update.", result);
#endif
                            k.SetValue(null, guid);
                            changed = true;
                            k.Close();
                        }
                    }
                }
            }
            catch (Exception e)
            {
                Common.OutputError(string.Format("Create regkey failed, '{0}'", result), e);
            }
            return changed;
        }
        internal static bool UnregisterKey(string ext, string key, ProviderType type)
        {
            bool changed = false;
            string result = ext + @"\shellex\" + key;
            try
            {
                switch (type)
                {
                    case ProviderType.ContextMenu:
                    case ProviderType.PropertySheet:
                        using (RegistryKey k = Registry.ClassesRoot.OpenSubKey(ext))
                        {
                            if (k != null)
                            {
                                result = k.GetValue(null) + @"\shellex\" + key;
                            }
                        }
                        break;
                }
                if (Registry.ClassesRoot.OpenSubKey(result) != null)
                {
#if DEBUG
                    Common.OutputDebugString("'{0}' - Delete key.", result);
#endif
                    Registry.ClassesRoot.DeleteSubKeyTree(result);
                    changed = true;
                }
            }
            catch (Exception e)
            {
                Common.OutputError(string.Format("Delete regkey failed, '{0}'", result), e);
            }
            return changed;
        }
    }

    static class Extension
    {
        public static IEnumerable<T> Where<T>(IEnumerable<T> list, Predicate<T> pred)
        {
            List<T> lst = new List<T>();
            if (list != null && pred != null)
            {
                foreach (T t in list)
                {
                    if (pred(t))
                    {
                        lst.Add(t);
                    }
                }
            }
            return lst;
        }
    }
}
